#pragma once
#include<iostream>
#include<algorithm>
using namespace std;

enum Colour
{
  RED,
  BLACK
};

template<class T>
struct RBTreeNode
{
  T _data;
  RBTreeNode<T>* _left;
  RBTreeNode<T>* _right;
  RBTreeNode<T>* _parent;

  Colour _col;

  RBTreeNode(const T& data)
    :_data(data)
     , _left(nullptr)
     , _right(nullptr)
     , _parent(nullptr)
  {}
};

template<class T, class Ref, class Ptr>
struct RBTreeIterator
{
  typedef RBTreeNode<T> Node;
  typedef RBTreeIterator<T, Ref, Ptr> Self;
  Node* _node;
  Node* _root;

  RBTreeIterator(Node* node, Node* root)
    :_node(node)
    ,_root(root)
  {}

  Ref operator*()
  {
    return _node->_data;
  }

  Ptr operator->()
  {
    return &_node->_data;
  }

  Self& operator++()
  {
    if(_node->_right)
    {
      Node* minleft = _node->_right;
      while(minleft->_left)
      {
        minleft = minleft->_left;
      }
      _node = minleft;
    }
    else{
      Node* cur = _node;
      Node* parent = cur->_parent;
      while(parent && cur == parent->_right)
      {
        cur = parent;
        parent = parent->_parent;
      }
      _node = parent;
    }
    return *this;
  }

  Self& operator--()
  {
    if(_node == nullptr)
    {
      Node* rightMost = _root;
      while(rightMost && rightMost->_right)
      {
        rightMost = rightMost->_right;
      }
      _node = rightMost;
    }
    else if(_node->_left)
    {
      Node* rightMost = _node->_left;
      while(rightMost->_right)
      {
        rightMost = rightMost->_right;
      }
      _node = rightMost;
    }
    else{
      Node* cur = _node;
      Node* parent = cur->_parent;
      while(parent && cur == parent->_left)
      {
        cur = parent;
        parent = cur->_parent;
      }
      _node = parent;
    }
    return *this;
  }

  bool operator!= (const Self& s)
  {
    return _node != s._node;
  }

  bool operator==(const Self& s)
  {
    return _node == s._node;
  }
};

template<class K, class T, class KeyOfT>
class RBTree
{
  typedef RBTreeNode<T> Node;
public:
  typedef RBTreeIterator<T, T&, T*> Iterator;
  typedef RBTreeIterator<T, const T&, const T*> ConstIterator;

  Iterator Begin()
  {
    Node* cur = _root;
    while(cur && cur->_left)
    {
      cur = cur->_left;
    }
    return Iterator(cur, _root);
  }

  Iterator End()
  {
    return Iterator(nullptr, _root);
  }

  ConstIterator Begin() const{
    Node* cur = _root;
    while(cur && cur->_left)
    {
      cur = cur->_left;
    }
    return ConstIterator(cur, _root);
  }

  ConstIterator End() const
  {
    return ConstIterator(nullptr, _root);
  }

  pair<Iterator, bool> Insert(const T& data)
  {
    if(_root == nullptr)
    {
      _root = new Node(data);
      _root->_col = BLACK;
      return make_pair(Iterator(_root, _root), true);
    }
    
    KeyOfT kot;
    Node* parent = nullptr;
    Node* cur = _root;
    while(cur)
    {
      if(kot(cur->_data) < kot(data))
      {
        parent = cur;
        cur = cur->_right;
      }
      else if(kot(cur->_data) > kot(data))
      {
        parent = cur;
        cur = cur->_left;
      }
      else{
        return make_pair(Iterator(cur, _root), false);
      }
    }
    cur = new Node(data);
    Node* newnode = cur;
    cur->_col = RED;
    if(kot(parent->_data) < kot(data))
    {
      parent->_right = cur;
    }
    else{
      parent->_left = cur;
    }
    cur->_parent = parent;
    while(parent && parent->_col == RED)
    {
      Node* grandfather = parent->_parent;
      if(grandfather->_left == parent)
      {
        Node* uncle = grandfather->_right;
        if(uncle && uncle->_col == RED)
        {
          //变色
          parent->_col = uncle->_col = BLACK;
          grandfather->_col = RED;

          //继续向上处理
          cur = grandfather;
          parent = cur->_parent;
        }
        else//uncle不存在，或者存在且为黑
        {
          if(cur == parent->_left)
          {
            //旋转+变色
            //    g
            //  p   u
            //c
            RotateR(grandfather);
            parent->_col = BLACK;
            grandfather->_col = RED;
          }
          else{
            //旋转+ 变色
            //   g
            // p  u
            //  c
            RotateL(parent);
            RotateR(grandfather);
            cur->_col = BLACK;
            grandfather->_col = RED;
          }
          break;
        }
      }
      else{
        //   g
        // u  p
        Node* uncle = grandfather->_left;
        //叔叔存在且为红
        if(uncle && uncle->_col == RED)
        {
          parent->_col = uncle->_col = BLACK;
          grandfather->_col = RED;
          //继续往上处理
          cur = grandfather;
          parent = cur->_parent;
        }
        else{//叔叔存在且为黑或不存在
          //    g
          //  u   p
          //       c
          if(cur == parent->_right)
          {
            RotateL(grandfather);
            parent->_col = BLACK;
            grandfather->_col = RED;
          }
          else{
            //    g
            //  u   p
            //    c
            RotateR(parent);
            RotateL(grandfather);
            cur->_col = BLACK;
            grandfather->_col = RED;
          }
          break;
        }
      }
    }
    _root->_col = BLACK;
    return make_pair(Iterator(newnode, _root), false);
  }
  

  Iterator Find(const K& key)
  {
    Node* cur = _root;
    while(cur)
    {
      if(cur->_kv.first < key)
      {
        cur = cur->_right;
      }
      else if(cur->_kv.first > key)
      {
        cur = cur->_left;
      }
      else{
        return Iterator(cur, _root);
      }
    }
    return End();
  }

protected:
  void RotateR(Node* parent)
  {
    Node* subL = parent->_left;
    Node* subLR = subL->_right;

    parent->_left = subLR;
    if(subLR)
      subLR->_parent = parent;
    Node* ppNode = parent->_parent;
    subL->_right = parent;
    parent->_parent = subL;

    if(parent == _root)
    {
      _root = subL;
      subL->_parent = nullptr;
    }
    else{
      if(ppNode->_left == parent)
      {
        ppNode->_left = subL;
      }
      else{
        ppNode->_right = subL;
      }
      subL->_parent = ppNode;
    }
  }

  void RotateL(Node* parent)
  {
    Node* subR = parent->_right;
    Node* subRL = subR->_left;
    parent->_right = subRL;
    if(subRL) 
      subRL->_parent = parent;
    Node* parentParent = parent->_parent;
    subR->_left = parent;
    parent->_parent = subR;

    if(parentParent == nullptr)
    {
      _root = subR;
      subR->_parent = nullptr;
    }
    else{
      if(parent == parentParent->_left)
      {
        parentParent->_left = subR;
      }
      else{
        parentParent->_right = subR;
      }
      subR->_parent = parentParent;
    }
  }


private:
  Node* _root = nullptr;
};

